import numpy as np

from .unit import Unit


class CharacterIndex(Unit):
    """
    CharacterIndexUnit for DIIN model.

    The input of :class:'CharacterIndexUnit' should be a list of word
    character list extracted from a text. The output is the character
    index representation of this text.

    :class:`NgramLetterUnit` and :class:`VocabularyUnit` are two
    essential prerequisite of :class:`CharacterIndexUnit`.

    Examples:
        >>> input_ = [['#', 'a', '#'],['#', 'o', 'n', 'e', '#']]
        >>> character_index = CharacterIndex(
        ...     char_index={
        ...      '<PAD>': 0, '<OOV>': 1, 'a': 2, 'n': 3, 'e':4, '#':5},
        ...     fixed_length_text=2,
        ...     fixed_length_word=5)
        >>> index = character_index.transform(input_)
        >>> index
        [[5.0, 2.0, 5.0, 0.0, 0.0], [5.0, 1.0, 3.0, 4.0, 5.0]]

    """

    def __init__(
        self,
        char_index: dict,
        fixed_length_text: int,
        fixed_length_word: int
    ):
        """
        Class initialization.

        :param char_index: character-index mapping generated by
            :class:'VocabularyUnit'.
        :param fixed_length_text: maximize length of a text.
        :param fixed_length_word: maximize length of a word.
        """
        self._char_index = char_index
        self._fixed_length_text = fixed_length_text
        self._fixed_length_word = fixed_length_word

    def transform(self, input_: list) -> list:
        """
        Transform list of characters to corresponding indices.

        :param input_: list of characters generated by
            :class:'NgramLetterUnit'.

        :return: character index representation of a text.
        """
        idx = np.zeros((self._fixed_length_text, self._fixed_length_word))
        for i in range(min(len(input_), self._fixed_length_text)):
            for j in range(min(len(input_[i]), self._fixed_length_word)):
                idx[i, j] = self._char_index.get(input_[i][j], 1)

        return idx.tolist()
